// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2011, 2013 Jitse Niesen <jitse@maths.leeds.ac.uk>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.

#ifndef EIGEN_MATRIX_SQUARE_ROOT
#define EIGEN_MATRIX_SQUARE_ROOT

namespace Eigen {

namespace internal {

    // pre:  T.block(i,i,2,2) has complex conjugate eigenvalues
    // post: sqrtT.block(i,i,2,2) is square root of T.block(i,i,2,2)
    template <typename MatrixType, typename ResultType> void matrix_sqrt_quasi_triangular_2x2_diagonal_block(const MatrixType& T, Index i, ResultType& sqrtT)
    {
        // TODO: This case (2-by-2 blocks with complex conjugate eigenvalues) is probably hidden somewhere
        //       in EigenSolver. If we expose it, we could call it directly from here.
        typedef typename traits<MatrixType>::Scalar Scalar;
        Matrix<Scalar, 2, 2> block = T.template block<2, 2>(i, i);
        EigenSolver<Matrix<Scalar, 2, 2>> es(block);
        sqrtT.template block<2, 2>(i, i) = (es.eigenvectors() * es.eigenvalues().cwiseSqrt().asDiagonal() * es.eigenvectors().inverse()).real();
    }

    // pre:  block structure of T is such that (i,j) is a 1x1 block,
    //       all blocks of sqrtT to left of and below (i,j) are correct
    // post: sqrtT(i,j) has the correct value
    template <typename MatrixType, typename ResultType>
    void matrix_sqrt_quasi_triangular_1x1_off_diagonal_block(const MatrixType& T, Index i, Index j, ResultType& sqrtT)
    {
        typedef typename traits<MatrixType>::Scalar Scalar;
        Scalar tmp = (sqrtT.row(i).segment(i + 1, j - i - 1) * sqrtT.col(j).segment(i + 1, j - i - 1)).value();
        sqrtT.coeffRef(i, j) = (T.coeff(i, j) - tmp) / (sqrtT.coeff(i, i) + sqrtT.coeff(j, j));
    }

    // similar to compute1x1offDiagonalBlock()
    template <typename MatrixType, typename ResultType>
    void matrix_sqrt_quasi_triangular_1x2_off_diagonal_block(const MatrixType& T, Index i, Index j, ResultType& sqrtT)
    {
        typedef typename traits<MatrixType>::Scalar Scalar;
        Matrix<Scalar, 1, 2> rhs = T.template block<1, 2>(i, j);
        if (j - i > 1)
            rhs -= sqrtT.block(i, i + 1, 1, j - i - 1) * sqrtT.block(i + 1, j, j - i - 1, 2);
        Matrix<Scalar, 2, 2> A = sqrtT.coeff(i, i) * Matrix<Scalar, 2, 2>::Identity();
        A += sqrtT.template block<2, 2>(j, j).transpose();
        sqrtT.template block<1, 2>(i, j).transpose() = A.fullPivLu().solve(rhs.transpose());
    }

    // similar to compute1x1offDiagonalBlock()
    template <typename MatrixType, typename ResultType>
    void matrix_sqrt_quasi_triangular_2x1_off_diagonal_block(const MatrixType& T, Index i, Index j, ResultType& sqrtT)
    {
        typedef typename traits<MatrixType>::Scalar Scalar;
        Matrix<Scalar, 2, 1> rhs = T.template block<2, 1>(i, j);
        if (j - i > 2)
            rhs -= sqrtT.block(i, i + 2, 2, j - i - 2) * sqrtT.block(i + 2, j, j - i - 2, 1);
        Matrix<Scalar, 2, 2> A = sqrtT.coeff(j, j) * Matrix<Scalar, 2, 2>::Identity();
        A += sqrtT.template block<2, 2>(i, i);
        sqrtT.template block<2, 1>(i, j) = A.fullPivLu().solve(rhs);
    }

    // solves the equation A X + X B = C where all matrices are 2-by-2
    template <typename MatrixType>
    void matrix_sqrt_quasi_triangular_solve_auxiliary_equation(MatrixType& X, const MatrixType& A, const MatrixType& B, const MatrixType& C)
    {
        typedef typename traits<MatrixType>::Scalar Scalar;
        Matrix<Scalar, 4, 4> coeffMatrix = Matrix<Scalar, 4, 4>::Zero();
        coeffMatrix.coeffRef(0, 0) = A.coeff(0, 0) + B.coeff(0, 0);
        coeffMatrix.coeffRef(1, 1) = A.coeff(0, 0) + B.coeff(1, 1);
        coeffMatrix.coeffRef(2, 2) = A.coeff(1, 1) + B.coeff(0, 0);
        coeffMatrix.coeffRef(3, 3) = A.coeff(1, 1) + B.coeff(1, 1);
        coeffMatrix.coeffRef(0, 1) = B.coeff(1, 0);
        coeffMatrix.coeffRef(0, 2) = A.coeff(0, 1);
        coeffMatrix.coeffRef(1, 0) = B.coeff(0, 1);
        coeffMatrix.coeffRef(1, 3) = A.coeff(0, 1);
        coeffMatrix.coeffRef(2, 0) = A.coeff(1, 0);
        coeffMatrix.coeffRef(2, 3) = B.coeff(1, 0);
        coeffMatrix.coeffRef(3, 1) = A.coeff(1, 0);
        coeffMatrix.coeffRef(3, 2) = B.coeff(0, 1);

        Matrix<Scalar, 4, 1> rhs;
        rhs.coeffRef(0) = C.coeff(0, 0);
        rhs.coeffRef(1) = C.coeff(0, 1);
        rhs.coeffRef(2) = C.coeff(1, 0);
        rhs.coeffRef(3) = C.coeff(1, 1);

        Matrix<Scalar, 4, 1> result;
        result = coeffMatrix.fullPivLu().solve(rhs);

        X.coeffRef(0, 0) = result.coeff(0);
        X.coeffRef(0, 1) = result.coeff(1);
        X.coeffRef(1, 0) = result.coeff(2);
        X.coeffRef(1, 1) = result.coeff(3);
    }

    // similar to compute1x1offDiagonalBlock()
    template <typename MatrixType, typename ResultType>
    void matrix_sqrt_quasi_triangular_2x2_off_diagonal_block(const MatrixType& T, Index i, Index j, ResultType& sqrtT)
    {
        typedef typename traits<MatrixType>::Scalar Scalar;
        Matrix<Scalar, 2, 2> A = sqrtT.template block<2, 2>(i, i);
        Matrix<Scalar, 2, 2> B = sqrtT.template block<2, 2>(j, j);
        Matrix<Scalar, 2, 2> C = T.template block<2, 2>(i, j);
        if (j - i > 2)
            C -= sqrtT.block(i, i + 2, 2, j - i - 2) * sqrtT.block(i + 2, j, j - i - 2, 2);
        Matrix<Scalar, 2, 2> X;
        matrix_sqrt_quasi_triangular_solve_auxiliary_equation(X, A, B, C);
        sqrtT.template block<2, 2>(i, j) = X;
    }

    // pre:  T is quasi-upper-triangular and sqrtT is a zero matrix of the same size
    // post: the diagonal blocks of sqrtT are the square roots of the diagonal blocks of T
    template <typename MatrixType, typename ResultType> void matrix_sqrt_quasi_triangular_diagonal(const MatrixType& T, ResultType& sqrtT)
    {
        using std::sqrt;
        const Index size = T.rows();
        for (Index i = 0; i < size; i++)
        {
            if (i == size - 1 || T.coeff(i + 1, i) == 0)
            {
                eigen_assert(T(i, i) >= 0);
                sqrtT.coeffRef(i, i) = sqrt(T.coeff(i, i));
            }
            else
            {
                matrix_sqrt_quasi_triangular_2x2_diagonal_block(T, i, sqrtT);
                ++i;
            }
        }
    }

    // pre:  T is quasi-upper-triangular and diagonal blocks of sqrtT are square root of diagonal blocks of T.
    // post: sqrtT is the square root of T.
    template <typename MatrixType, typename ResultType> void matrix_sqrt_quasi_triangular_off_diagonal(const MatrixType& T, ResultType& sqrtT)
    {
        const Index size = T.rows();
        for (Index j = 1; j < size; j++)
        {
            if (T.coeff(j, j - 1) != 0)  // if T(j-1:j, j-1:j) is a 2-by-2 block
                continue;
            for (Index i = j - 1; i >= 0; i--)
            {
                if (i > 0 && T.coeff(i, i - 1) != 0)  // if T(i-1:i, i-1:i) is a 2-by-2 block
                    continue;
                bool iBlockIs2x2 = (i < size - 1) && (T.coeff(i + 1, i) != 0);
                bool jBlockIs2x2 = (j < size - 1) && (T.coeff(j + 1, j) != 0);
                if (iBlockIs2x2 && jBlockIs2x2)
                    matrix_sqrt_quasi_triangular_2x2_off_diagonal_block(T, i, j, sqrtT);
                else if (iBlockIs2x2 && !jBlockIs2x2)
                    matrix_sqrt_quasi_triangular_2x1_off_diagonal_block(T, i, j, sqrtT);
                else if (!iBlockIs2x2 && jBlockIs2x2)
                    matrix_sqrt_quasi_triangular_1x2_off_diagonal_block(T, i, j, sqrtT);
                else if (!iBlockIs2x2 && !jBlockIs2x2)
                    matrix_sqrt_quasi_triangular_1x1_off_diagonal_block(T, i, j, sqrtT);
            }
        }
    }

}  // end of namespace internal

/** \ingroup MatrixFunctions_Module
  * \brief Compute matrix square root of quasi-triangular matrix.
  *
  * \tparam  MatrixType  type of \p arg, the argument of matrix square root,
  *                      expected to be an instantiation of the Matrix class template.
  * \tparam  ResultType  type of \p result, where result is to be stored.
  * \param[in]  arg      argument of matrix square root.
  * \param[out] result   matrix square root of upper Hessenberg part of \p arg.
  *
  * This function computes the square root of the upper quasi-triangular matrix stored in the upper
  * Hessenberg part of \p arg.  Only the upper Hessenberg part of \p result is updated, the rest is
  * not touched.  See MatrixBase::sqrt() for details on how this computation is implemented.
  *
  * \sa MatrixSquareRoot, MatrixSquareRootQuasiTriangular
  */
template <typename MatrixType, typename ResultType> void matrix_sqrt_quasi_triangular(const MatrixType& arg, ResultType& result)
{
    eigen_assert(arg.rows() == arg.cols());
    result.resize(arg.rows(), arg.cols());
    internal::matrix_sqrt_quasi_triangular_diagonal(arg, result);
    internal::matrix_sqrt_quasi_triangular_off_diagonal(arg, result);
}

/** \ingroup MatrixFunctions_Module
  * \brief Compute matrix square root of triangular matrix.
  *
  * \tparam  MatrixType  type of \p arg, the argument of matrix square root,
  *                      expected to be an instantiation of the Matrix class template.
  * \tparam  ResultType  type of \p result, where result is to be stored.
  * \param[in]  arg      argument of matrix square root.
  * \param[out] result   matrix square root of upper triangular part of \p arg.
  *
  * Only the upper triangular part (including the diagonal) of \p result is updated, the rest is not
  * touched.  See MatrixBase::sqrt() for details on how this computation is implemented.
  *
  * \sa MatrixSquareRoot, MatrixSquareRootQuasiTriangular
  */
template <typename MatrixType, typename ResultType> void matrix_sqrt_triangular(const MatrixType& arg, ResultType& result)
{
    using std::sqrt;
    typedef typename MatrixType::Scalar Scalar;

    eigen_assert(arg.rows() == arg.cols());

    // Compute square root of arg and store it in upper triangular part of result
    // This uses that the square root of triangular matrices can be computed directly.
    result.resize(arg.rows(), arg.cols());
    for (Index i = 0; i < arg.rows(); i++) { result.coeffRef(i, i) = sqrt(arg.coeff(i, i)); }
    for (Index j = 1; j < arg.cols(); j++)
    {
        for (Index i = j - 1; i >= 0; i--)
        {
            // if i = j-1, then segment has length 0 so tmp = 0
            Scalar tmp = (result.row(i).segment(i + 1, j - i - 1) * result.col(j).segment(i + 1, j - i - 1)).value();
            // denominator may be zero if original matrix is singular
            result.coeffRef(i, j) = (arg.coeff(i, j) - tmp) / (result.coeff(i, i) + result.coeff(j, j));
        }
    }
}

namespace internal {

    /** \ingroup MatrixFunctions_Module
  * \brief Helper struct for computing matrix square roots of general matrices.
  * \tparam  MatrixType  type of the argument of the matrix square root,
  *                      expected to be an instantiation of the Matrix class template.
  *
  * \sa MatrixSquareRootTriangular, MatrixSquareRootQuasiTriangular, MatrixBase::sqrt()
  */
    template <typename MatrixType, int IsComplex = NumTraits<typename internal::traits<MatrixType>::Scalar>::IsComplex> struct matrix_sqrt_compute
    {
        /** \brief Compute the matrix square root
    *
    * \param[in]  arg     matrix whose square root is to be computed.
    * \param[out] result  square root of \p arg.
    *
    * See MatrixBase::sqrt() for details on how this computation is implemented.
    */
        template <typename ResultType> static void run(const MatrixType& arg, ResultType& result);
    };

    // ********** Partial specialization for real matrices **********

    template <typename MatrixType> struct matrix_sqrt_compute<MatrixType, 0>
    {
        typedef typename MatrixType::PlainObject PlainType;
        template <typename ResultType> static void run(const MatrixType& arg, ResultType& result)
        {
            eigen_assert(arg.rows() == arg.cols());

            // Compute Schur decomposition of arg
            const RealSchur<PlainType> schurOfA(arg);
            const PlainType& T = schurOfA.matrixT();
            const PlainType& U = schurOfA.matrixU();

            // Compute square root of T
            PlainType sqrtT = PlainType::Zero(arg.rows(), arg.cols());
            matrix_sqrt_quasi_triangular(T, sqrtT);

            // Compute square root of arg
            result = U * sqrtT * U.adjoint();
        }
    };

    // ********** Partial specialization for complex matrices **********

    template <typename MatrixType> struct matrix_sqrt_compute<MatrixType, 1>
    {
        typedef typename MatrixType::PlainObject PlainType;
        template <typename ResultType> static void run(const MatrixType& arg, ResultType& result)
        {
            eigen_assert(arg.rows() == arg.cols());

            // Compute Schur decomposition of arg
            const ComplexSchur<PlainType> schurOfA(arg);
            const PlainType& T = schurOfA.matrixT();
            const PlainType& U = schurOfA.matrixU();

            // Compute square root of T
            PlainType sqrtT;
            matrix_sqrt_triangular(T, sqrtT);

            // Compute square root of arg
            result = U * (sqrtT.template triangularView<Upper>() * U.adjoint());
        }
    };

}  // end namespace internal

/** \ingroup MatrixFunctions_Module
  *
  * \brief Proxy for the matrix square root of some matrix (expression).
  *
  * \tparam Derived  Type of the argument to the matrix square root.
  *
  * This class holds the argument to the matrix square root until it
  * is assigned or evaluated for some other reason (so the argument
  * should not be changed in the meantime). It is the return type of
  * MatrixBase::sqrt() and most of the time this is the only way it is
  * used.
  */
template <typename Derived> class MatrixSquareRootReturnValue : public ReturnByValue<MatrixSquareRootReturnValue<Derived>>
{
protected:
    typedef typename internal::ref_selector<Derived>::type DerivedNested;

public:
    /** \brief Constructor.
      *
      * \param[in]  src  %Matrix (expression) forming the argument of the
      * matrix square root.
      */
    explicit MatrixSquareRootReturnValue(const Derived& src) : m_src(src) {}

    /** \brief Compute the matrix square root.
      *
      * \param[out]  result  the matrix square root of \p src in the
      * constructor.
      */
    template <typename ResultType> inline void evalTo(ResultType& result) const
    {
        typedef typename internal::nested_eval<Derived, 10>::type DerivedEvalType;
        typedef typename internal::remove_all<DerivedEvalType>::type DerivedEvalTypeClean;
        DerivedEvalType tmp(m_src);
        internal::matrix_sqrt_compute<DerivedEvalTypeClean>::run(tmp, result);
    }

    Index rows() const { return m_src.rows(); }
    Index cols() const { return m_src.cols(); }

protected:
    const DerivedNested m_src;
};

namespace internal {
    template <typename Derived> struct traits<MatrixSquareRootReturnValue<Derived>>
    {
        typedef typename Derived::PlainObject ReturnType;
    };
}  // namespace internal

template <typename Derived> const MatrixSquareRootReturnValue<Derived> MatrixBase<Derived>::sqrt() const
{
    eigen_assert(rows() == cols());
    return MatrixSquareRootReturnValue<Derived>(derived());
}

}  // end namespace Eigen

#endif  // EIGEN_MATRIX_FUNCTION
